import sqlite3
from . import custommodule
from datetime import datetime
import pandas as pd

# 1、显示自选股列表
def display_watchlist(number:int=1):
    watchlist = custommodule.get_watchlist(number)
    if not watchlist:
        print("\n您的自选股列表为空")
        return
    
    print("\n您的自选股列表：")
    print("=" * 40)
    print(f"{'代码':<10} {'名称'}")
    print("-" * 40)
    for stock_code, stock_name in watchlist:
        print(f"{stock_code:<10} {stock_name}")
    print("=" * 40)
    print(f"共 {len(watchlist)} 只股票")

# 2、添加单个自选股
def add_to_watchlist(query:str=None):
    if query is None:
        query = input("\n请输入要添加的股票代码或名称：").strip()

    # 搜索匹配的股票
    results = custommodule.search_stock(query)
    
    if not results:
        print(f"⚠️ 未找到与 '{query}' 相关的股票，请先更新股票列表")
        return
    
    # 如果只有一个结果，直接使用
    if len(results) == 1:
        stock_code, stock_name = results[0]
        print(f"找到股票：{stock_code} - {stock_name}")
    else:
        print("\n找到多个相关股票：")
        for i, (code, name) in enumerate(results):
            print(f"{i+1}. {code} - {name}")
        add_to_watchlist()
        return  # 添加return语句，防止继续执行后面的代码
    
    conn = sqlite3.connect('db_file.db')
    cursor = conn.cursor()
    
    # 检查是否已经是自选股
    cursor.execute("SELECT is_watchlist FROM stock_info WHERE stock_code = ?", (stock_code,))
    is_watchlist = cursor.fetchone()[0]
    
    if is_watchlist == 1:
        print(f"⚠️ 股票 {stock_code} 已经是您的自选股")
        conn.close()
        return
    
    # 更新为自选股
    cursor.execute("UPDATE stock_info SET is_watchlist = 1 WHERE stock_code = ?", (stock_code,))
    conn.commit()
    conn.close()
    print(f"✅ 股票 {stock_code} ({stock_name}) 成功添加到自选股列表")

    # 全量更新该股数据
    custommodule.update_stock_daily([stock_code])
    custommodule.calculate_technical_indicators([stock_code])


# 3、删除单个自选股
def remove_from_watchlist(query:str=None):
    if query is None:
        query = input("\n请输入要删除的股票代码或名称：").strip()
    
    # 搜索匹配的股票
    results = custommodule.search_stock(query)
    
    if not results:
        print(f"⚠️ 未找到与 '{query}' 相关的股票")
        return
    
    # 如果只有一个结果，直接使用
    if len(results) == 1:
        stock_code, stock_name = results[0]
        print(f"找到股票：{stock_code} - {stock_name}")
    # 如果有多个结果，让用户选择
    else:
        print("\n找到多个相关股票：")
        for i, (code, name) in enumerate(results):
            print(f"{i+1}. {code} - {name}")
        remove_from_watchlist()
        return  # 添加return语句，防止继续执行后面的代码
    
    conn = sqlite3.connect('db_file.db')
    cursor = conn.cursor()
    
    # 检查股票是否在自选股列表中
    cursor.execute("SELECT is_watchlist FROM stock_info WHERE stock_code = ?", (stock_code,))
    is_watchlist = cursor.fetchone()[0]
    
    if not is_watchlist:
        print(f"⚠️ 股票 {stock_code} 不在您的自选股列表中")
        conn.close()
        return
    
    # 移除自选股状态
    cursor.execute("UPDATE stock_info SET is_watchlist = 0 WHERE stock_code = ?", (stock_code,))
    conn.commit()
    conn.close()
    print(f"✅ 股票 {stock_code} ({stock_name}) 已从自选股列表中移除")

# 4、自选股诊断
def analyze_watchlist(stock_codes:list=None, trade_date:str=None):
    """
    对自选股进行简单诊断分析
    """
    # 处理股票代码列表
    if stock_codes is None:
        stock_codes = custommodule.get_watchlist()
        if not stock_codes:
            return "您要诊断的股票列表为空"
        stock_codes = [code for code, _ in stock_codes]  # 获取自选股代码列表
    stock_codes = custommodule.StockCodeConversion.remove(stock_codes)  # 去掉代码后缀
    # 处理交易日期
    if trade_date is None:
        with sqlite3.connect('db_file.db') as conn:
            cursor = conn.cursor()
            trade_date = datetime.now().strftime("%Y-%m-%d")
            if trade_date not in cursor.execute("SELECT trade_date FROM stock_daily").fetchone()[0]:
                trade_date = cursor.execute("SELECT MAX(trade_date) FROM stock_daily").fetchone()[0]
    # 处理查询指标
    stock_indicators = ["Angle_State"]

    # 生成查询字段（股票代码、名称 + 用户选的指标）
    select_fields = "stock_daily.stock_code, stock_info.stock_name, stock_daily.trade_date"
    for i in stock_indicators:
        select_fields = select_fields + f", stock_technical.{i}"

    # 查询日线及指标数据
    with sqlite3.connect('db_file.db') as conn:

        # 动态生成WHERE子句的占位符
        if len(stock_codes) == 1:
            where_condition = "stock_daily.stock_code = ?"
        else:
            placeholders = ", ".join(["?"] * len(stock_codes))  # 生成多个占位符（如"?, ?"）
            where_condition = f"stock_daily.stock_code IN ({placeholders})"

        # 构建查询语句（使用动态生成的where_condition）
        query = f"""
                    SELECT {select_fields}
                    FROM stock_daily
                    INNER JOIN stock_technical ON stock_daily.stock_code = stock_technical.stock_code AND stock_daily.trade_date = stock_technical.trade_date
                    INNER JOIN stock_info ON stock_info.stock_code LIKE (stock_daily.stock_code || '.%')
                    WHERE {where_condition} AND stock_daily.trade_date = ?
                """

        # 准备参数：股票代码列表 + 交易日期
        params = stock_codes + [trade_date]
        result = pd.read_sql(query, conn, params=params)

    # 计算操作建议
    operate = custommodule.Strategy(stock_codes=stock_codes, start_date=trade_date, end_date=trade_date).result
    # 合并查询结果和操作建议
    if operate is None:
        results = result
    else:
        results = pd.merge(result, operate, on=['stock_code', 'trade_date'], how='inner')

    return results

# 自选股管理主菜单
def menu():
    while True:
        print("\n" + "=" * 20)
        print(" 自选股管理菜单 ")
        print("=" * 20)
        print("1. 查看自选股")
        print("2. 添加自选股")
        print("3. 删除自选股")
        print("4. 查看推荐股")
        print("0. 返回主菜单")
        print("-" * 20)
        choice = input("请输入选项（0-4）：").strip()
        
        # 处理用户选择
        match choice:
            case '1':
                display_watchlist(1)
            case '2':
                add_to_watchlist()
            case '3':
                remove_from_watchlist()
            case '4':
                display_watchlist(2)
            case '0':
                break
            case _:
                print("⚠️ 输入错误，请输入 0-4 之间的数字")
